import tushare as ts
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
# 设置tushare的token
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()
# 选择一支股票
stock_code = '600520.SH'
# 获取3 - 5年的行情数据
start_date = '20220101'
end_date = '20250101'
df = pro.daily(ts_code=stock_code, start_date=start_date, end_date=end_date)
# 按日期升序排序
df = df.sort_values(by='trade_date')
df['trade_date'] = pd.to_datetime(df['trade_date'])
df.set_index('trade_date', inplace=True)
# 计算技术指标与特征值
df['pct_change'] = df['close'].pct_change()
df['ma5'] = df['close'].rolling(window=5).mean()
df['ma10'] = df['close'].rolling(window=10).mean()
df['ma20'] = df['close'].rolling(window=20).mean()
df['std5'] = df['close'].rolling(window=5).std()
df['std10'] = df['close'].rolling(window=10).std()
df['std20'] = df['close'].rolling(window=20).std()
# 去除缺失值
df = df.dropna()
# 定义特征和目标变量
X = df[['open', 'high', 'low', 'close', 'pre_close', 'ma5', 'ma10', 'ma20', 'std5', 'std10', 'std20']]
y = df['pct_change'].shift(-1).dropna()
X = X[:-1]
# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# 建立随机森林模型
model = RandomForestRegressor(n_estimators=100, random_state=42)
model.fit(X_train, y_train)
# 进行预测
y_pred = model.predict(X_test)
# 计算收益率
df['predicted_pct_change'] = np.nan
df.loc[df.index[-len(y_pred):], 'predicted_pct_change'] = y_pred
df['strategy_return'] = (1 + df['predicted_pct_change']).cumprod()
# 模型评价
mse = mean_squared_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)
print(f"均方误差 (MSE): {mse}")
print(f"决定系数 (R^2): {r2}")